package com.abc.service.impl;

import com.abc.dto.request.ChatRequestDto;
import com.abc.dto.request.KnowledgeRequestDto;
import com.abc.dto.response.ChatResponseDto;
import com.abc.dto.response.KnowledgeResponseDto;
import com.abc.service.ChatService;
import com.abc.service.SessionService;
import com.abc.tool.ToolService;
import com.abc.util.ValidationUtil;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import javax.validation.Valid;
import java.util.Collections;
import java.util.Map;

@Slf4j
@Service
public class ChatServiceImpl implements ChatService {

    @Autowired
    private ChatClient dashScopeChatClient;

    @Autowired
    private SessionService sessionService;

    @Autowired
    private RetrievalAugmentationAdvisor retrievalAugmentationAdvisor;

    @Autowired
    private ToolService toolService;


    @Override
    public Flux<ChatResponseDto> send(@Valid ChatRequestDto chatRequestDto) {
        return ValidationUtil
                .validateChatRequest(sessionService, chatRequestDto.getSessionId(), chatRequestDto.getMessage())
                .thenMany(
                        dashScopeChatClient
                                .prompt(chatRequestDto.getMessage())
                                .advisors(params -> params.param(ChatMemory.CONVERSATION_ID, chatRequestDto.getSessionId()))
                                .tools(toolService)
                                .stream()
                                .content()
                                .map(ChatResponseDto::new)
                );
    }

    @Override
    public Flux<KnowledgeResponseDto> knowledge(KnowledgeRequestDto chatRequestDto) {

        return ValidationUtil.validateChatRequest(sessionService, chatRequestDto.getSessionId(), chatRequestDto.getMessage())
                .thenMany(
                        dashScopeChatClient
                                .prompt()
                                .system("你也可以结合工具解答用户的问题")
                                .user(chatRequestDto.getMessage())
                                .advisors(params -> params.param(ChatMemory.CONVERSATION_ID, chatRequestDto.getSessionId()))
                                .advisors(retrievalAugmentationAdvisor)
                                .advisors(new SimpleLoggerAdvisor())
                                .stream()
                                .chatResponse()
                                .map(r -> {
                                    KnowledgeResponseDto response = new KnowledgeResponseDto();
                                    response.setMessage(r.getResult().getOutput().getText());
                                    response.setDocuments(Collections.singletonList(JSON.toJSONString(r.getResult())));
                                    log.info("大模型返回结果: {}", JSON.toJSONString(response));
                                    return response;
                                })

                );
    }



}
